import os, shutil
from glob import glob

import numpy as np
from PIL import Image, ImageFile
ImageFile.LOAD_TRUNCATED_IMAGES = True

from keras.models import load_model
from tqdm import tqdm

import pickle, json, fire, time

with open('config.json', 'r', encoding='utf-8') as f:
    config = json.load(f)

document_clf = load_model(config['DOCUMENT_CLF']['MODEL'], compile=False)
layout_clf = load_model(config['LAYOUT_CLF']['MODEL'], compile=False)

with open(config['DOCUMENT_CLF']['LABELS'], 'rb') as document_label_f:
    document_label = pickle.load(document_label_f)

with open(config['LAYOUT_CLF']['LABELS'], 'rb') as layout_label_f:
    layout_label = pickle.load(layout_label_f)

def main(ismove=False,endtime=None):
    sourcePath = config['sourcePath']
    if not os.path.exists(sourcePath):
        print('源路径不存在')
        return
    nowtime = time.strftime('%Y%m%d', time.localtime(time.time()))
    if endtime is not None:nowtime=endtime
    ForeCastTime = config['ForeCastTime']

    daysList = glob(os.path.join(sourcePath, '*'))
    for day in daysList:
        if ForeCastTime is None or int(ForeCastTime) > int(os.path.split(day)[1]) < int(nowtime): continue
        ## 证件类型采样
        DocumentClassifierDatasets = os.path.join(day, 'DocumentClassifier\datasets')
        find_files(DocumentClassifierDatasets, 'document_error_{}.txt'.format(nowtime), document_clf, document_label,
                   ismove)
        ## 版面类型采样
        LayoutClassifierDatasets = os.path.join(day, 'LayoutClassifier\datasets')
        find_files(LayoutClassifierDatasets, 'layout_error_{}.txt'.format(nowtime), layout_clf, layout_label, ismove)
        # 采样完成后记录采样时间
        config['ForeCastTime'] = os.path.split(day)[1]
    f.close()

    with open('config.json', 'w', encoding='utf-8') as outfile:
        outfile.write(json.dumps(config, sort_keys=False, indent=2, separators=(',', ':'), ensure_ascii=False))
        outfile.close()

'''迭代文件夹 'document_error.txt'  'DocumentClassifier\\datasets'  '''
def find_files(datasets, logtxt, clf, label, ismove=False):
    if os.path.exists(datasets):
        with open(logtxt, 'a') as log:
            for filename_ in os.listdir(datasets):
                fp_ = os.path.join(datasets, filename_)
                datasetsFiles = glob(os.path.join(fp_, '*'))
                for i in tqdm(range(len(datasetsFiles)), desc=fp_):
                    src = datasetsFiles[i]
                    predictType = predictOne(src, clf, label)
                    if predictType != filename_:
                        log.write('{}\t{}\t{}\n'.format(src, predictType, filename_))
                        targetPath = os.path.join(datasets.replace('原图', '错误'), filename_, predictType,
                                                  os.path.split(src)[1])
                        if ismove: movefile(src, targetPath)
        log.close()
    pass


'''对图片进行预测'''


def predictOne(input_img, clf, labels=None):
    input_shape = clf.get_input_shape_at(0)
    if isinstance(input_img, str):
        img = Image.open(input_img)
    elif isinstance(input_img, Image.Image):
        img = input_img
    else:
        raise TypeError('classifierImage input_img is path string or pillow image object')
    img_arr_formated = formatImage(img, input_shape[1:])
    X = img_arr_formated.astype('float32')
    X = 1 - X / 255
    X = X.reshape(-1, input_shape[1], input_shape[2], input_shape[3])

    probs = clf.predict(X)[0]
    predict = (probs == probs.max()).dot(np.arange(probs.size))
    return labels[int(predict)]


'''将图片转成数组'''


def formatImage(image, input_shape):
    src_size = np.array(image.size)
    dst_size = np.array([input_shape[0], input_shape[1]])

    if image.mode != 'RGB': image = image.convert('RGB')
    trans_array = np.zeros(input_shape, dtype='uint8') + 255

    k = dst_size / src_size.max()
    scale_size = np.around((k * src_size)).astype('int')
    trans_vect = ((dst_size - scale_size) / 2).astype('int')

    scale_image = image.resize(tuple(scale_size), resample=Image.BOX)
    scale_array = np.array(scale_image, dtype='uint8')
    trans_array[trans_vect[1]: trans_vect[1] + scale_size[1],
    trans_vect[0]: trans_vect[0] + scale_size[0]] = \
        np.array(scale_image, dtype='uint8')

    return trans_array


'''移动图片到指定目录'''


def movefile(srcfile, dstfile):
    if not os.path.isfile(srcfile):
        print("%s not exist!" % (srcfile))
    else:
        fpath, fname = os.path.split(dstfile)  # 分离文件名和路径
        if not os.path.exists(fpath):
            os.makedirs(fpath)  # 创建路径
        shutil.move(srcfile, dstfile)  # 移动文件
        print("move %s -> %s" % (srcfile, dstfile))


if __name__ == '__main__':
    fire.Fire(main)
